from langchain.document_loaders import TextLoader
from langchain.embeddings import OpenAIEmbeddings
from langchain.text_splitter import CharacterTextSplitter
from langchain.vectorstores.faiss import FAISS
from langchain.vectorstores.qdrant import Qdrant


def normal_vector_storage():
    raw_documents = TextLoader('../document_transformers/files/state_of_the_union.txt').load()
    text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=0)
    documents = text_splitter.split_documents(raw_documents)
    db = FAISS.from_documents(documents, OpenAIEmbeddings())
    query = "What did the president say about Ketanji Brown Jackson"
    docs = db.similarity_search(query)
    print(docs[0].page_content)
    embedding_vector = OpenAIEmbeddings().embed_query(query)
    docs = db.similarity_search_by_vector(embedding_vector)
    print(docs[0].page_content)


async def async_vector_storage():
    raw_documents = TextLoader('../document_transformers/files/state_of_the_union.txt').load()
    embeddings = OpenAIEmbeddings()
    db = await Qdrant.afrom_documents(raw_documents, embeddings, "http://localhost:6333")
    query = "What did the president say about Ketanji Brown Jackson"
    docs = await db.asimilarity_search(query)
    print(docs[0].page_content)
    embedding_vector = embeddings.embed_query(query)
    docs = await db.asimilarity_search_by_vector(embedding_vector)
    print(docs)


if __name__ == '__main__':
    normal_vector_storage()